package com.metal.gear.demo;

import com.alibaba.fastjson.JSON;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.PropertyNamingStrategy;
import com.metal.gear.ClientV4;
import com.metal.gear.Constants;
import com.metal.gear.service.v4.embedding.EmbeddingApiResponse;
import com.metal.gear.service.v4.embedding.EmbeddingRequest;
import com.metal.gear.service.v4.file.FileApiResponse;
import com.metal.gear.service.v4.file.QueryFileApiResponse;
import com.metal.gear.service.v4.file.QueryFilesRequest;
import com.metal.gear.service.v4.fine_turning.*;
import com.metal.gear.service.v4.image.CreateImageRequest;
import com.metal.gear.service.v4.image.ImageApiResponse;
import com.metal.gear.service.v4.model.*;
import com.metal.gear.service.v4.fine_turning.*;
import com.metal.gear.service.v4.model.*;
import io.reactivex.Flowable;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;


public class V4OkHttpClientTest {

    //private static final String API_KEY = "a58ebb3dd2ac093736beb9bcac869da2.Fm2sDbacg2Mkongp";
    //private static final String API_KEY = "4a0bb594391a9546d63f2d07152f5da1.DR7D283qmmjxSWRS";
    private static final String API_KEY = "c2e55ef666f32ffc2372c92abff1f467.vvnXw89O6VIJI5dz";


    private static final ClientV4 client = new ClientV4.Builder(API_KEY).build();

    private static final ObjectMapper mapper = defaultObjectMapper();


    public static ObjectMapper defaultObjectMapper() {
        ObjectMapper mapper = new ObjectMapper();
        mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
        mapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
        mapper.setPropertyNamingStrategy(PropertyNamingStrategy.SNAKE_CASE);
        mapper.addMixIn(ChatFunction.class, ChatFunctionMixIn.class);
        mapper.addMixIn(ChatCompletionRequest.class, ChatCompletionRequestMixIn.class);
        mapper.addMixIn(ChatFunctionCall.class, ChatFunctionCallMixIn.class);
        return mapper;
    }

    // 请自定义自己的业务id
    private static final String requestIdTemplate = "wpp-first-%d";

    public static void test() throws Exception {
        testInvoke();
    };

    public static void main(String[] args) throws Exception {
        //System.setProperty("org.slf4j.simpleLogger.logFile", "System.out");
        // 1. sse-invoke调用模型，使用标准Listener，直接返回结果
        //testSseInvoke();

        // 2. invoke调用模型,直接返回结果
        testInvoke();

        // 3. 异步调用
//         String taskId = testAsyncInvoke();
        // 4.异步查询
//         testQueryResult(taskId);

        // 5.文生图
          //testCreateImage();

        // 6. 图生文
         // testImageToWord();

        // 7. 向量模型
        // testEmbeddings();

        // 8.微调-上传微调数据集
         // testUploadFile();

        // 9.微调-查询上传文件列表
          //testQueryUploadFileList();

        // 10.微调-创建微调任务
         //testCreateFineTuningJob();

        // 11.微调-查询微调任务事件
//          testQueryFineTuningJobsEvents();

        // 12.微调-查询微调任务
//        testRetrieveFineTuningJobs();

        // 13.微调-查询个人微调任务
//          testQueryPersonalFineTuningJobs();

        // 14.微调-调用微调模型（参考模型调用接口，并替换成要调用模型的编码model）
    }

    private static void testQueryPersonalFineTuningJobs() {
        QueryPersonalFineTuningJobRequest queryPersonalFineTuningJobRequest = new QueryPersonalFineTuningJobRequest();
        queryPersonalFineTuningJobRequest.setLimit(1);
        QueryPersonalFineTuningJobApiResponse queryPersonalFineTuningJobApiResponse = client.queryPersonalFineTuningJobs(queryPersonalFineTuningJobRequest);
        System.out.println("model output:" + JSON.toJSONString(queryPersonalFineTuningJobApiResponse));

    }

    private static void testQueryFineTuningJobsEvents() {
        QueryFineTuningJobRequest queryFineTuningJobRequest = new QueryFineTuningJobRequest();
        queryFineTuningJobRequest.setJobId("ftjob-20240119114544390-zkgjb");
//        queryFineTuningJobRequest.setLimit(1);
//        queryFineTuningJobRequest.setAfter("1");
        QueryFineTuningEventApiResponse queryFineTuningEventApiResponse = client.queryFineTuningJobsEvents(queryFineTuningJobRequest);
        System.out.println("model output:" + JSON.toJSONString(queryFineTuningEventApiResponse));
    }

    /**
     * 查询微调任务
     */
    private static void testRetrieveFineTuningJobs() {
        QueryFineTuningJobRequest queryFineTuningJobRequest = new QueryFineTuningJobRequest();
        queryFineTuningJobRequest.setJobId("ftjob-20240119114544390-zkgjb");
//        queryFineTuningJobRequest.setLimit(1);
//        queryFineTuningJobRequest.setAfter("1");
        QueryFineTuningJobApiResponse queryFineTuningJobApiResponse = client.retrieveFineTuningJobs(queryFineTuningJobRequest);
        System.out.println("model output:" + JSON.toJSONString(queryFineTuningJobApiResponse));
    }

    /**
     * 创建微调任务
     */
    private static void testCreateFineTuningJob() {
        FineTuningJobRequest request = new FineTuningJobRequest();
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        request.setRequestId(requestId);
        //可以重写命名
        request.setModel("chatglm3-6b");
        //上传微调文件之后返回的名称id
        request.setTraining_file("file-20240506093803964-vbxz8");
        CreateFineTuningJobApiResponse createFineTuningJobApiResponse = client.createFineTuningJob(request);
        System.out.println("model output:" + JSON.toJSONString(createFineTuningJobApiResponse));
    }

    /**
     * 微调文件上传列表查询
     */
    private static void testQueryUploadFileList() {
        QueryFilesRequest queryFilesRequest = new QueryFilesRequest();
        QueryFileApiResponse queryFileApiResponse = client.queryFilesApi(queryFilesRequest);
        System.out.println("model output:" + JSON.toJSONString(queryFileApiResponse));
    }


    /**
     * 微调上传数据集
     */
    private static void testUploadFile() {
        String filePath = "/Users/wangdelian/Documents/idea_work/大模型/zhipuai-sdk-java-v4/file/transaction-data.jsonl";
        String purpose = "fine-tune";
        FileApiResponse fileApiResponse = client.invokeUploadFileApi(purpose, filePath);
        System.out.println("model output:" + JSON.toJSONString(fileApiResponse));
    }

    private static void testEmbeddings() {
        EmbeddingRequest embeddingRequest = new EmbeddingRequest();
        //设置输入的文本信息
        embeddingRequest.setInput("hello world");
        //设置模型  embedding-2
        embeddingRequest.setModel(Constants.ModelEmbedding2);
        //模型调用
        EmbeddingApiResponse apiResponse = client.invokeEmbeddingsApi(embeddingRequest);
        //结果输出
        System.out.println("model output:" + JSON.toJSONString(apiResponse));
    }

    /**
     * 图生文
     */
    private static void testImageToWord() {
        List<ChatMessage> messages = new ArrayList<>();
        List<Map<String, Object>> contentList = new ArrayList<>();
        //请求的问题描述
        Map<String, Object> textMap = new HashMap<>();
        textMap.put("type", "text");
        textMap.put("text", "图里有什么,看起来健康不健康");
        //图片地址
        Map<String, Object> typeMap = new HashMap<>();
        typeMap.put("type", "image_url");
        Map<String, Object> urlMap = new HashMap<>();
        urlMap.put("url", "https://sfile.chatglm.cn/testpath/275ae5b6-5390-51ca-a81a-60332d1a7cac_0.png");
        typeMap.put("image_url", urlMap);
        contentList.add(textMap);
        contentList.add(typeMap);
        //请求体
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), contentList);
        messages.add(chatMessage);
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());

        //模型调用
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                //设置模型 glm-4v 通用大模型
                .model(Constants.ModelChatGLM4V)
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .requestId(requestId)
                .build();
        ModelApiResponse modelApiResponse = client.invokeModelApi(chatCompletionRequest);
        System.out.println("model output:" + JSON.toJSONString(modelApiResponse));

    }

    private static void testCreateImage() {
        CreateImageRequest createImageRequest = new CreateImageRequest();
        //设置图像模型  cogview
        createImageRequest.setModel(Constants.ModelCogView);
        //请求描述
        createImageRequest.setPrompt("画一个温顺可爱的小狗");
        //模型调用
        ImageApiResponse imageApiResponse = client.createImage(createImageRequest);
        //结果返回
        System.out.println("imageApiResponse:" + JSON.toJSONString(imageApiResponse));
    }


    /**
     * sse调用
     */
    private static void testSseInvoke() {
        List<ChatMessage> messages = new ArrayList<>();
        //请求问题描述
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), "写一个java的分布式锁");
        ChatMessage chatMessage1 = new ChatMessage(ChatMessageRole.ASSISTANT.value(), "分布式锁是一种机制，用于在分布式系统中保持数据的同步和一致性，确保在多个并发进程或服务器中，同一时间只有一个任务可以操作特定的资源。\n" +
                "\n" +
                "以下是使用Java实现一个简单的分布式锁的示例，这个例子使用了Redis作为分布式锁的后端存储。我们将使用Jedis库来与Redis进行交互。\n" +
                "\n" +
                "首先，你需要将Jedis依赖项添加到你的项目中。如果你使用Maven，可以在`pom.xml`中添加以下依赖：\n" +
                "\n" +
                "```xml\n" +
                "<dependencies>\n" +
                "    <!-- Jedis: Java client for Redis -->\n" +
                "    <dependency>\n" +
                "        <groupId>redis.clients</groupId>\n" +
                "        <artifactId>jedis</artifactId>\n" +
                "        <version>3.7.0</version> <!-- 请使用最新的版本号 -->\n" +
                "    </dependency>\n" +
                "</dependencies>\n" +
                "```\n" +
                "\n" +
                "以下是分布式锁的简单实现：\n" +
                "\n" +
                "```java\n" +
                "import redis.clients.jedis.Jedis;\n" +
                "import redis.clients.jedis.params.SetParams;\n" +
                "\n" +
                "public class DistributedLock {\n" +
                "\n" +
                "    private Jedis jedis;\n" +
                "    private String lockKey;\n" +
                "    private String lockValue;\n" +
                "    private int timeout = 30; // 锁的超时时间（秒）\n" +
                "\n" +
                "    public DistributedLock(Jedis jedis, String lockKey) {\n" +
                "        this.jedis = jedis;\n" +
                "        this.lockKey = lockKey;\n" +
                "        this.lockValue = String.valueOf(System.currentTimeMillis());\n" +
                "    }\n" +
                "\n" +
                "    /**\n" +
                "     * 尝试获取锁\n" +
                "     * @return true if the lock is acquired, false otherwise\n" +
                "     */\n" +
                "    public boolean tryLock() {\n" +
                "        SetParams params = new SetParams();\n" +
                "        params.nx().px(timeout * 1000); // nx表示只有当key不存在时才设置，px设置超时时间（毫秒）\n" +
                "\n" +
                "        String result = jedis.set(lockKey, lockValue, params);\n" +
                "        return \"OK\".equals(result);\n" +
                "    }\n" +
                "\n" +
                "    /**\n" +
                "     * 释放锁\n" +
                "     * @return true if the lock is released successfully, false otherwise\n" +
                "     */\n" +
                "    public boolean releaseLock() {\n" +
                "        // Lua script to release the lock safely\n" +
                "        String script = \"if redis.call('get', KEYS[1]) == ARGV[1] then \" +\n" +
                "                \"return redis.call('del', KEYS[1]) else return 0 end\";\n" +
                "        \n" +
                "        Object result = jedis.eval(script, 1, lockKey, lockValue);\n" +
                "        \n" +
                "        return \"1\".equals(result.toString());\n" +
                "    }\n" +
                "}\n" +
                "\n" +
                "```\n" +
                "\n" +
                "以下是使用这个分布式锁的例子：\n" +
                "\n" +
                "```java\n" +
                "public class Main {\n" +
                "    public static void main(String[] args) {\n" +
                "        // 连接到Redis\n" +
                "        Jedis jedis = new Jedis(\"localhost\", 6379);\n" +
                "        String lockKey = \"myLock\"; // 锁的标识\n" +
                "\n" +
                "        // 创建分布式锁\n" +
                "        DistributedLock lock = new DistributedLock(jedis, lockKey);\n" +
                "\n" +
                "        // 尝试获取锁\n" +
                "        if (lock.tryLock()) {\n" +
                "            try {\n" +
                "                // 执行业务逻辑\n" +
                "                System.out.println(\"Lock acquired, executing business logic...\");\n" +
                "                // 模拟业务逻辑耗时\n" +
                "                Thread.sleep(2000);\n" +
                "            } catch (InterruptedException e) {\n" +
                "                e.printStackTrace();\n" +
                "            } finally {\n" +
                "                // 释放锁\n" +
                "                if (lock.releaseLock()) {\n" +
                "                    System.out.println(\"Lock released successfully.\");\n" +
                "                } else {\n" +
                "                    System.out.println(\"Failed to release the lock.\");\n" +
                "                }\n" +
                "            }\n" +
                "        } else {\n" +
                "            System.out.println(\"Failed to acquire the lock.\");\n" +
                "        }\n" +
                "\n" +
                "        // 关闭Jedis连接\n" +
                "        jedis.close();\n" +
                "    }\n" +
                "}\n" +
                "```\n" +
                "\n" +
                "请注意，这是一个简单的例子，并没有涵盖所有分布式锁的复杂性，例如重试策略、死锁处理、锁续期等。在实际应用中，你可能需要使用更加成熟的解决方案，比如基于Redis的Redlock算法或者使用Zookeeper等。");
        ChatMessage chatMessage2 = new ChatMessage(ChatMessageRole.USER.value(), "再优化一下");
       // ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), "上一个问题有什么可以优化的地方");
        messages.add(chatMessage);
        messages.add(chatMessage1);
        messages.add(chatMessage2);
        //设置请求id
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        //封装一个请求对象
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                //指定使用模型
                .model(Constants.ModelChatGLM4)
                //是否异步返回
                .stream(Boolean.TRUE)
                //请求问题描述
                .messages(messages)
                //请求唯一id
                .requestId(requestId)
                .toolChoice("auto")
                .build();
        // 模型调用   apikey调用方式
        // 登录注册免费的API Key    private static final String API_KEY = "a58ebb3dd2ac093736beb9bcac869da2.Fm2sDbacg2Mkongp";
        // new一个客户端直接使用     private static final ClientV4 client = new ClientV4.Builder(API_KEY).build();
        ModelApiResponse sseModelApiResp = client.invokeModelApi(chatCompletionRequest);
        //结果输出
        if (sseModelApiResp.isSuccess()) {
            AtomicBoolean isFirst = new AtomicBoolean(true);
            ChatMessageAccumulator chatMessageAccumulator = mapStreamToAccumulator(sseModelApiResp.getFlowable())
                    .doOnNext(accumulator -> {
                        {
                            if (isFirst.getAndSet(false)) {
                                System.out.print("Response: ");
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getTool_calls() != null) {
                                String jsonString = mapper.writeValueAsString(accumulator.getDelta().getTool_calls());
                                System.out.println("tool_calls: " + jsonString);
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getContent() != null) {
                                System.out.print(accumulator.getDelta().getContent());
                            }
                        }
                    })
                    .doOnComplete(System.out::println)
                    .lastElement()
                    .blockingGet();

            Choice choice = new Choice(chatMessageAccumulator.getChoice().getFinishReason(), 0L, chatMessageAccumulator.getDelta());
            List<Choice> choices = new ArrayList<>();
            choices.add(choice);
            ModelData data = new ModelData();
            data.setChoices(choices);
            data.setUsage(chatMessageAccumulator.getUsage());
            data.setId(chatMessageAccumulator.getId());
            data.setCreated(chatMessageAccumulator.getCreated());
            data.setRequestId(chatCompletionRequest.getRequestId());
            sseModelApiResp.setFlowable(null);
            sseModelApiResp.setData(data);
        }
        System.out.println("model output:" + JSON.toJSONString(sseModelApiResp));
    }

    public static Flowable<ChatMessageAccumulator> mapStreamToAccumulator(Flowable<ModelData> flowable) {
        return flowable.map(chunk -> {
            return new ChatMessageAccumulator(chunk.getChoices().get(0).getDelta(), null, chunk.getChoices().get(0), chunk.getUsage(), chunk.getCreated(), chunk.getId());
        });
    }

    /**
     * 同步调用
     */
    private static void testInvoke() {
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), "ChatGLM和你哪个更强大");
        messages.add(chatMessage);
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        // 函数调用参数构建部分
        List<ChatTool> chatToolList = new ArrayList<>();
        ChatTool chatTool = new ChatTool();
        chatTool.setType(ChatToolType.FUNCTION.value());
        ChatFunctionParameters chatFunctionParameters = new ChatFunctionParameters();
        chatFunctionParameters.setType("object");
        Map<String, Object> properties = new HashMap<>();
        properties.put("location", new HashMap<String, Object>() {{
            put("type", "string");
            put("description", "城市，如：北京");
        }});
        properties.put("unit", new HashMap<String, Object>() {{
            put("type", "string");
            put("enum", new ArrayList<String>() {{
                add("celsius");
                add("fahrenheit");
            }});
        }});
        chatFunctionParameters.setProperties(properties);
        ChatFunction chatFunction = ChatFunction.builder()
                .name("get_weather")
                .description("Get the current weather of a location")
                .parameters(chatFunctionParameters)
                .build();
        chatTool.setFunction(chatFunction);
        chatToolList.add(chatTool);
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                //.model(Constants.ModelChatGLM4)
                .model("chatglm3-6b-8617287520837443787")
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .requestId(requestId)
                .tools(chatToolList)
                .toolChoice("auto")
                .build();
        ModelApiResponse invokeModelApiResp = client.invokeModelApi(chatCompletionRequest);
        try {
            System.out.println("model output:" + mapper.writeValueAsString(invokeModelApiResp));
        } catch (JsonProcessingException e) {
            e.printStackTrace();
        }
    }


    /**
     * 异步调用
     */
    private static String testAsyncInvoke() {
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), "ChatLM和你哪个更强大");
        messages.add(chatMessage);
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        // 函数调用参数构建部分
        List<ChatTool> chatToolList = new ArrayList<>();
        ChatTool chatTool = new ChatTool();
        chatTool.setType(ChatToolType.FUNCTION.value());
        ChatFunctionParameters chatFunctionParameters = new ChatFunctionParameters();
        chatFunctionParameters.setType("object");
        Map<String, Object> properties = new HashMap<>();
        properties.put("location", new HashMap<String, Object>() {{
            put("type", "string");
            put("description", "城市，如：北京");
        }});
        properties.put("unit", new HashMap<String, Object>() {{
            put("type", "string");
            put("enum", new ArrayList<String>() {{
                add("celsius");
                add("fahrenheit");
            }});
        }});
        chatFunctionParameters.setProperties(properties);
        ChatFunction chatFunction = ChatFunction.builder()
                .name("get_weather")
                .description("Get the current weather of a location")
                .parameters(chatFunctionParameters)
                .build();
        chatTool.setFunction(chatFunction);
        chatToolList.add(chatTool);
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethodAsync)
                .messages(messages)
                .requestId(requestId)
                .tools(chatToolList)
                .toolChoice("auto")
                .build();
        ModelApiResponse invokeModelApiResp = client.invokeModelApi(chatCompletionRequest);
        System.out.println("model output:" + JSON.toJSONString(invokeModelApiResp));
        return invokeModelApiResp.getData().getTaskId();
    }


    /**
     * 查询异步结果
     *
     * @param taskId
     */
    private static void testQueryResult(String taskId) {
        QueryModelResultRequest request = new QueryModelResultRequest();
        request.setTaskId(taskId);
        QueryModelResultResponse queryResultResp = client.queryModelResult(request);
        try {
            System.out.println("model output:" + mapper.writeValueAsString(queryResultResp));
        } catch (JsonProcessingException e) {
            e.printStackTrace();
        }
    }
}

